from abc import ABC, abstractmethod
from typing import Dict, Any, Optional
from app.core.config import settings
import docker
import json
import subprocess
import asyncio
import logging

logger = logging.getLogger(__name__)


class BaseFuzzEngine(ABC):
    """模糊测试引擎的抽象基类"""
    
    def __init__(self, engine_name: str):
        self.engine_name = engine_name
    
    @abstractmethod
    async def initialize(self, config: Dict[str, Any]):
        """初始化引擎"""
        pass
    
    @abstractmethod
    async def execute(self, target_id: int, update_progress=None):
        """执行模糊测试"""
        pass
    
    @abstractmethod
    async def stop(self):
        """停止模糊测试"""
        pass
    
    @abstractmethod
    async def get_results(self):
        """获取测试结果"""
        pass


class DockerFuzzEngine(BaseFuzzEngine):
    """基于Docker的模糊测试引擎"""
    
    def __init__(self, engine_name: str):
        super().__init__(engine_name)
        self.client = docker.from_env()
        self.container = None
    
    async def initialize(self, config: Dict[str, Any]):
        """初始化Docker环境"""
        self.config = config
        logger.info(f"Initializing Docker engine {self.engine_name} with config: {config}")
    
    async def execute(self, target_id: int, update_progress=None):
        """在Docker容器中执行模糊测试"""
        try:
            # 这里是模拟执行，实际应用中应该根据不同引擎构建适当的容器
            image = f"fuzznexus/{self.engine_name}:latest"
            
            # 准备容器环境
            env = {
                "TARGET_ID": str(target_id),
                "CONFIG": json.dumps(self.config)
            }
            
            # 运行容器
            self.container = self.client.containers.run(
                image=image,
                environment=env,
                detach=True,
                # 限制资源
                mem_limit=self.config.get("memory_limit", "2g"),
                cpu_quota=self.config.get("cpu_quota", 100000),
                cpu_period=100000,
            )
            
            # 模拟进度更新
            if update_progress:
                for progress in range(0, 101, 10):
                    await asyncio.sleep(2)  # 模拟测试过程
                    update_progress(state='PROGRESS', meta={'progress': progress})
            
            # 等待容器完成
            self.container.wait()
            
            # 获取日志
            logs = self.container.logs().decode('utf-8')
            logger.info(f"Engine {self.engine_name} logs: {logs}")
            
            return {"status": "success", "logs": logs}
            
        except Exception as e:
            logger.error(f"Docker engine {self.engine_name} failed: {str(e)}")
            raise
        
    async def stop(self):
        """停止Docker容器"""
        if self.container:
            try:
                self.container.stop()
                self.container.remove()
                self.container = None
            except Exception as e:
                logger.error(f"Failed to stop container: {str(e)}")
    
    async def get_results(self):
        """获取容器中的结果"""
        # 这里应该实现从容器中复制结果文件的逻辑
        return {"results": []}


class NativeFuzzEngine(BaseFuzzEngine):
    """在本地环境执行的模糊测试引擎"""
    
    def __init__(self, engine_name: str):
        super().__init__(engine_name)
        self.process = None
    
    async def initialize(self, config: Dict[str, Any]):
        """初始化本地环境"""
        self.config = config
        logger.info(f"Initializing native engine {self.engine_name} with config: {config}")
    
    async def execute(self, target_id: int, update_progress=None):
        """在本地执行模糊测试"""
        try:
            # 这里是模拟执行，实际应用中应该调用相应的引擎命令
            cmd = [self.engine_name, "--target", str(target_id)]
            
            # 添加配置参数
            for key, value in self.config.items():
                cmd.extend([f"--{key}", str(value)])
            
            # 启动进程
            self.process = await asyncio.create_subprocess_exec(
                *cmd,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE
            )
            
            # 模拟进度更新
            if update_progress:
                for progress in range(0, 101, 10):
                    await asyncio.sleep(2)  # 模拟测试过程
                    update_progress(state='PROGRESS', meta={'progress': progress})
            
            # 等待进程完成
            stdout, stderr = await self.process.communicate()
            
            return {
                "status": "success",
                "stdout": stdout.decode('utf-8'),
                "stderr": stderr.decode('utf-8')
            }
            
        except Exception as e:
            logger.error(f"Native engine {self.engine_name} failed: {str(e)}")
            raise
    
    async def stop(self):
        """停止本地进程"""
        if self.process and self.process.returncode is None:
            self.process.terminate()
            try:
                await asyncio.wait_for(self.process.wait(), timeout=5.0)
            except asyncio.TimeoutError:
                self.process.kill()
    
    async def get_results(self):
        """获取本地结果"""
        # 这里应该实现从本地目录读取结果的逻辑
        return {"results": []}


class EngineManager:
    """引擎管理器，负责管理和调度不同的模糊测试引擎"""
    
    def __init__(self):
        self.engines = {
            "afl++": DockerFuzzEngine("afl++"),
            "libfuzzer": NativeFuzzEngine("libfuzzer"),
            "honggfuzz": DockerFuzzEngine("honggfuzz"),
            "american_fuzzy_lop": DockerFuzzEngine("afl")
        }
    
    async def execute_fuzz_task(
        self, 
        engine_name: str, 
        target_id: int, 
        config: Dict[str, Any],
        task_id: int,
        update_progress=None
    ) -> Dict[str, Any]:
        """执行模糊测试任务"""
        # 检查引擎是否支持
        if engine_name not in self.engines:
            raise ValueError(f"Unsupported engine: {engine_name}")
        
        engine = self.engines[engine_name]
        
        try:
            # 初始化引擎
            await engine.initialize(config)
            
            # 执行测试
            result = await engine.execute(target_id, update_progress)
            
            # 获取结果
            results = await engine.get_results()
            
            return {
                "task_id": task_id,
                "engine": engine_name,
                "execution_result": result,
                "fuzz_results": results
            }
            
        finally:
            # 确保停止引擎
            await engine.stop()
    
    def register_engine(self, engine_name: str, engine_class: BaseFuzzEngine):
        """注册新的引擎"""
        self.engines[engine_name] = engine_class(engine_name)
        logger.info(f"Registered new engine: {engine_name}")
    
    def list_engines(self) -> Dict[str, str]:
        """列出所有可用的引擎"""
        return {name: engine.__class__.__name__ for name, engine in self.engines.items()}